import os
import uuid
import shutil
import asyncio
import base64
from typing import Optional, List
from fastapi import FastAPI, HTTPException
from fastapi.responses import FileResponse, JSONResponse, Response
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
import uvicorn

# 导入图像生成函数
from async_comfiy_ui import generate_image_async

# 全局变量：控制图像生成状态
is_generating = False

app = FastAPI(title="ComfyUI图像生成API", description="使用ComfyUI生成图像的FastAPI接口")

# 配置CORS，允许前端调用
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 在生产环境中应该指定具体的域名
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

class ImageGenerationRequest(BaseModel):
    """图像生成请求模型"""
    positive_prompt: str = "一张美丽的风景画"
    negative_prompt: str = ""

@app.get("/generate_image")
async def generate_image_api(
    positive_prompt: str = "一张美丽的风景画",
    negative_prompt: str = "",
    width: int = 1024,
    height: int = 1024
):
    """
    生成图像的API接口
    
    Args:
        positive_prompt: 正面提示词
        negative_prompt: 负面提示词
        width: 图像宽度 (最大2048)
        height: 图像高度 (最大2048)
        
    Returns:
        JSONResponse: 包含base64编码图像的JSON响应
    """
    global is_generating
    
    # 验证图像尺寸
    if width > 2048 or height > 2048:
        raise HTTPException(
            status_code=400, 
            detail="图像尺寸超出限制：宽度和高度不能超过2048像素"
        )
    
    if width <= 0 or height <= 0:
        raise HTTPException(
            status_code=400, 
            detail="图像尺寸必须为正数"
        )
    
    # 检查是否有其他请求正在生成图像
    if is_generating:
        return JSONResponse(
            status_code=429,
            content={"message": "请等待，其他图片输出中", "status": "busy"}
        )
    
    try:
        # 设置生成状态为忙碌
        is_generating = True
        
        # 确保images文件夹存在
        os.makedirs("images", exist_ok=True)
        
        # 生成唯一的请求ID
        request_id = str(uuid.uuid4())
        
        print(f"开始生成图像 - 请求ID: {request_id}")
        print(f"提示词: {positive_prompt}")
        
        # 直接调用异步图像生成函数
        generated_images = await generate_image_async(
            positive_prompt=positive_prompt,
            negative_prompt=negative_prompt,
            width=width,
            height=height
        )
        
        if generated_images:
            # 处理第一张生成的图像
            for image_filename in generated_images:
                try:
                    # 检查文件是否存在
                    if os.path.exists(image_filename):
                        # 读取图像文件并转换为base64
                        with open(image_filename, "rb") as image_file:
                            image_data = image_file.read()
                            base64_image = base64.b64encode(image_data).decode('utf-8')
                            # 创建完整的Data URL格式
                            data_url = f"data:image/png;base64,{base64_image}"
                        
                        # 移动图像到images文件夹
                        filename_only = os.path.basename(image_filename)
                        target_path = os.path.join("images", filename_only)
                        shutil.move(image_filename, target_path)
                        
                        print(f"图像 {image_filename} 已移动到 {target_path}")
                        
                        # 返回包含base64图像的JSON响应
                        return JSONResponse(content={
                            "status": "success",
                            "message": "图像生成成功",
                            "image_data_url": data_url,
                            "filename": filename_only,
                            "request_id": request_id
                        })
                    else:
                        print(f"警告: 文件 {image_filename} 不存在")
                        continue
                except Exception as e:
                    print(f"处理图像 {image_filename} 时出错: {str(e)}")
                    continue
            
            # 如果没有成功处理任何图像
            raise HTTPException(status_code=500, detail="图像生成失败：无法处理生成的图像文件")
        else:
            raise HTTPException(status_code=500, detail="图像生成失败：没有生成任何图像")
            
    except Exception as e:
        print(f"图像生成错误: {str(e)}")
        raise HTTPException(status_code=500, detail=f"图像生成失败: {str(e)}")
    finally:
        # 无论成功还是失败，都要重置生成状态
        is_generating = False

@app.post("/generate_image")
async def generate_image_api_post(request: ImageGenerationRequest):
    """
    生成图像的API接口 (POST方法，返回base64编码的图像)
    
    Args:
        request: 图像生成请求参数
        
    Returns:
        JSONResponse: 包含base64编码图像的JSON响应
    """
    global is_generating
    
    # 验证图像尺寸
    if request.width > 2048 or request.height > 2048:
        raise HTTPException(
            status_code=400, 
            detail="图像尺寸超出限制：宽度和高度不能超过2048像素"
        )
    
    if request.width <= 0 or request.height <= 0:
        raise HTTPException(
            status_code=400, 
            detail="图像尺寸必须为正数"
        )
    
    # 检查是否有其他请求正在生成图像
    if is_generating:
        return JSONResponse(
            status_code=429,
            content={"message": "请等待，其他图片输出中", "status": "busy"}
        )
    
    try:
        # 设置生成状态为忙碌
        is_generating = True
        
        # 确保images文件夹存在
        os.makedirs("images", exist_ok=True)
        
        # 生成唯一的请求ID
        request_id = str(uuid.uuid4())
        
        print(f"开始生成图像 - 请求ID: {request_id}")
        print(f"提示词: {request.positive_prompt}")
        
        # 直接调用异步图像生成函数
        generated_images = await generate_image_async(
            positive_prompt=request.positive_prompt,
            negative_prompt=request.negative_prompt,
            width=request.width,
            height=request.height
        )
        
        if generated_images:
            # 处理第一张生成的图像
            for image_filename in generated_images:
                try:
                    # 检查文件是否存在
                    if os.path.exists(image_filename):
                        # 读取图像文件并转换为base64
                        with open(image_filename, "rb") as image_file:
                            image_data = image_file.read()
                            base64_image = base64.b64encode(image_data).decode('utf-8')
                            # 创建完整的Data URL格式
                            data_url = f"data:image/png;base64,{base64_image}"
                        
                        # 移动图像到images文件夹
                        filename_only = os.path.basename(image_filename)
                        target_path = os.path.join("images", filename_only)
                        shutil.move(image_filename, target_path)
                        
                        print(f"图像 {image_filename} 已移动到 {target_path}")
                        
                        # 返回包含base64图像的JSON响应
                        return JSONResponse(content={
                            "status": "success",
                            "message": "图像生成成功",
                            "image_data_url": data_url,
                            "filename": filename_only,
                            "request_id": request_id
                        })
                    else:
                        print(f"警告: 文件 {image_filename} 不存在")
                        continue
                except Exception as e:
                    print(f"处理图像 {image_filename} 时出错: {str(e)}")
                    continue
            
            # 如果没有成功处理任何图像
            raise HTTPException(status_code=500, detail="图像生成失败：无法处理生成的图像文件")
        else:
            raise HTTPException(status_code=500, detail="图像生成失败：没有生成任何图像")
            
    except Exception as e:
        print(f"图像生成错误: {str(e)}")
        raise HTTPException(status_code=500, detail=f"图像生成失败: {str(e)}")
    finally:
        # 无论成功还是失败，都要重置生成状态
        is_generating = False

@app.get("/status")
async def get_status():
    """
    获取API状态
    """
    global is_generating
    return JSONResponse(content={
        "is_generating": is_generating,
        "message": "API正在运行" if not is_generating else "有图像正在生成中"
    })

if __name__ == "__main__":
    print("启动ComfyUI图像生成API服务器...")
    print("API文档地址: http://localhost:8000/docs")
    print("ReDoc文档地址: http://localhost:8000/redoc")
    
    uvicorn.run(
        "app:app",
        host="0.0.0.0",
        port=8000,
        reload=True
    )